import os
import numpy as np
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import mpl_toolkits.mplot3d.axes3d as p3
from matplotlib.animation import FuncAnimation

from trajectory_compensate import check_traj_continuity

# plt.figure()
# plt.scatter(point_trajectory[:,0], point_trajectory[:,1], c='r')
# plt.show()

def draw_arrow(point_trajectory, start_points, directions):
    # 绘制轨迹和方向箭头
    plt.figure()
    plt.scatter(point_trajectory[:, 0], point_trajectory[:, 1], c='r', label='trajectory')
    plt.quiver(start_points[:, 0], start_points[:, 1], directions[:, 0], directions[:, 1], color='y', scale=1000, width=0.003, label='direction')
    plt.title('transform_trajectory')
    plt.xlabel('X coordinate')
    plt.ylabel('Y coordinate')
        # 调整坐标轴方向，使得左上角为原点
    plt.gca().invert_yaxis()
    plt.legend()
    plt.grid(True)
    plt.show()

def draw_trajectory_2d(traj: np.ndarray, jumps=None, save_gif=False, filename=None):
    # 设置坐标轴范围
    xMax, xMin = traj[:, 0].max(), traj[:, 0].min()
    yMax, yMin = traj[:, 1].max(), traj[:, 1].min()
    xPadding = 0.2 * (xMax - xMin)    
    yPadding = 0.2 * (yMax - yMin)

    
    fig, ax = plt.subplots(figsize=(8, 6))
    # plt.scatter(traj[:,0], traj[:,1], c='r')
    # 画散点图,每两点之间连线,连线末端加箭头
    # dxdy = np.diff(traj, axis=0)
    for i, p in enumerate(traj):
        plt.scatter(p[0], p[1], c='r')
        # if i < len(traj)-1:
            # ax.arrow(traj[i,0], traj[i,1], dxdy[i,0], dxdy[i,1], 
            #     head_width=2, head_length=5, fc='y', ec='y')
        if i > 0:
            plt.annotate('', 
                xy=(traj[i,0], traj[i,1]), xycoords='data',
                xytext=(traj[i-1,0], traj[i-1,1]), textcoords='data',
                arrowprops=dict(arrowstyle="->", color="goldenrod", lw=1, ls=':'))
    if jumps is not None:
        for pid in jumps:
            # 两点之间连线，连线末端加箭头
            # plt.plot([traj[pid-1,0], traj[pid,0]], [traj[pid-1,1], traj[pid,1]], 'r--')
            plt.annotate('', 
                    xy=(traj[pid,0], traj[pid,1]), xycoords='data',
                    xytext=(traj[pid-1,0], traj[pid-1,1]), textcoords='data',
                    arrowprops=dict(arrowstyle="->", color="gray", lw=2, ls=':'))
    # 绘制起点
    plt.scatter(traj[0, 0], traj[0, 1], color='green', marker='o', s=100, label='Start')
    # 绘制终点
    plt.scatter(traj[-1, 0], traj[-1, 1], color='k', marker='*', lw=3, s=100, label='End')
    
    # 添加图例
    plt.legend()
    # 设置坐标轴范围
    ax.set_xlim(xMin - xPadding, xMax + xPadding)
    ax.set_ylim(yMin - yPadding, yMax + yPadding)
    
    # 调整坐标轴方向，使得左上角为原点
    plt.gca().invert_yaxis()
    
    if save_gif:
        if not filename:
            filename = os.path.join(WORK_DIR, "trajectory_2D.jpg")
        # 保存JPG
        plt.savefig(filename, dpi=100)
        
    plt.show()


def draw_trajectory_2d_animation(point_trajectory, jumps=None, save_gif=False, filename=None):
    
    # 计算每个点到下一个点的方向向量
    directions = np.diff(point_trajectory, axis=0)

    # 为了在每个点上绘制箭头，我们需要在每个点上加上一个起点
    start_points = point_trajectory[:-1]

    # 创建图形和坐标轴
    fig, ax = plt.subplots()

    # 初始化散点和箭头
    ax.scatter(point_trajectory[0, 0], point_trajectory[0, 1], color='green', marker='o', s=100, label='Start', zorder=5)
    scatter = ax.scatter(start_points[0, 0], start_points[0, 1], c='r') # , label='trajectory'
    # quiver = ax.quiver(start_points[0, 0], start_points[0, 1], 
    #                    directions[0, 0], directions[0, 1], color='y', 
    #                    scale=500, width=0.003) # , label='direction'

    ax.scatter(point_trajectory[-1, 0], point_trajectory[-1, 1], color='k', lw=3, marker='*', s=100, label='End', zorder=5)
    # 设置坐标轴范围
    padding = 0.2 * (point_trajectory.max() - point_trajectory.min())
    ax.set_xlim(point_trajectory[:, 0].min() - padding, point_trajectory[:, 0].max() + padding)
    ax.set_ylim(point_trajectory[:, 1].min() - padding, point_trajectory[:, 1].max() + padding)

    
    ax.legend()
    ax.set_title('transform_trajectory')
    ax.set_xlabel('X (H)')
    ax.set_ylabel('Y (W)')
    ax.grid(True)
    ax.invert_yaxis()  # 添加此行以反转Y轴

    # 更新箭头位置的函数
    def update(frame):
        scatter.set_offsets(point_trajectory[:frame+1, :])  # 更新点位置    
        # quiver.set_offsets(start_points[frame, :])  # 更新箭头起点
        # quiver.set_UVC(directions[frame, 0], directions[frame, 1])  # 更新箭头方向
        
        if jumps is not None and frame in jumps:
            ax.scatter(point_trajectory[frame, 0], point_trajectory[frame, 1], c='gray', marker='p')
        # if frame > 0:
        #     print(f"frame: {frame}, dist: {np.linalg.norm(point_trajectory[frame] - point_trajectory[frame-1])}")
        # ax.scatter([start_points[:frame+1, 0]], [start_points[:frame+1, 1]], c='r', label='trajectory')
        # ax.quiver(start_points[:frame, 0], start_points[:frame, 1], directions[:frame, 0], directions[:frame, 1], color='y', scale=1000, width=0.003, label='direction')

    
    # 创建动画
    ani = FuncAnimation(fig, update, frames=len(directions), interval=100, repeat=False)

    # # 调整坐标轴方向，使得左上角为原点
    # plt.gca().invert_yaxis()
    
    if save_gif:
        if filename is None:
            # 保存动画为GIF, 标题增加时间戳
            import time
            timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
            filename = os.path.join(WORK_DIR, f'trajectory_2D_{timestamp}.gif')
        ani.save(filename, writer='pillow', fps=10)
        # ani.save('trajectory.gif', writer='pillow', fps=10)

    # 显示动画
    plt.show()


def draw_trajectory_3d(traj_3d: np.ndarray, save_fig=0, filename=None):
    """_summary_

    Args:
        traj_3d (np.ndarray): N x 3, 3D trajectory points
    """
    # # 获取 figure 画布并创建 Axes3D 对象
    fig = plt.figure(figsize=(12, 8))
    # ax = Axes3D(fig)
    # ax = fig.add_subplot(111, projection='3d')
    

    # 通过子图创建 Axes3D 对象
    ax: Axes3D = plt.subplot(111, projection='3d')
    
    # N x 3 -> 3 x N
    x, y, z = traj_3d.T
    
    # xMax, xMin = x.max(), x.min()
    # yMax, yMin = y.max(), y.min()
    # zMax, zMin = z.max(), z.min()

    # ax.scatter(0,0,0, c='b', marker='p') 
    ax.scatter(x[0], y[0], z[0], color='green', marker='o', s=100, label='Start')    
       
    ax.plot(x, y, z, c='r', marker='o', ls=':')
    
    # 绘制箭头，从起点指向终点
    for x0, y0, z0, x1, y1, z1 in zip(x[:-1], y[:-1], z[:-1], x[1:], y[1:], z[1:]):
        ax.quiver(x0, y0, z0, x1-x0, y1-y0, z1-z0, 
                color='y', length=0.5, arrow_length_ratio=0.3, normalize=0)
    
    # # Make the direction data for the arrows
    # u = np.sin(np.pi * x) * np.cos(np.pi * y) * np.cos(np.pi * z)
    # v = -np.cos(np.pi * x) * np.sin(np.pi * y) * np.cos(np.pi * z)
    # w = (np.sqrt(2.0 / 3.0) * np.cos(np.pi * x) * np.cos(np.pi * y) *
    #     np.sin(np.pi * z))

    # ax.quiver(x, y, z, u, v, w, length=1, normalize=0)
    
    ax.scatter(x[-1], y[-1], z[-1], color='k', marker='*', s=100, label='End')
    # ax.plot(x, y, z)

    ax.set_xlabel('X (W)')
    ax.set_ylabel('Y (D)')
    ax.set_zlabel('Z (H)')
    plt.title('3D trajectory')
    
    if save_fig:
        if not filename:
            filename = os.path.join(WORK_DIR, "trajectory_3D.jpg")
        # 保存JPG
        plt.savefig(filename, dpi=100)
        
    plt.show()

def cell_interpolate(frame, trajectory, inpl_space=None, **kwargs):
    point = trajectory[frame]
    
    yMin = kwargs.get('yMin', None)
    yMax = kwargs.get('yMax', None)
    zMin = kwargs.get('zMin', None)
    zMax = kwargs.get('zMax', None)
    
    dist = np.linalg.norm(point - trajectory[frame-1])        
    x, y, z = None, None, None
    
    if inpl_space is not None and dist > inpl_space: 
        
        if frame > 0:
            num_inpl = int(dist / inpl_space) + 2
            # print(f"frame: {frame}, dist: {dist}, inpl_space: {inpl_space}, num_inpl: {num_inpl}")
            
            if num_inpl < 3  or num_inpl > len(trajectory):
                print(f"frame: {frame}, num_inpl: {num_inpl}, dist: {dist}, inpl_space: {inpl_space}")
                return x, y, z
            
        if frame > 2 and frame < len(trajectory)-2:
            # 获取前一个点和当前点
            # p00 = trajectory[frame-3]
            p0 = trajectory[frame-2]
            p1 = trajectory[frame-1]
            p2 = point
            p3 = trajectory[frame+1]
            p4 = trajectory[frame+2]
            
            pp = np.vstack((p0, p1, p2, p3, p4))
            
            if len(point) == 3: 
                x, y, z = pp.T
                y = np.vstack((y, z)).T
            else:
                x, y = pp.T
            
            # 使用 np.polyfit 拟合一个二次函数
            coefficients = np.polyfit(x, y, 2)
            
            nx = np.linspace(p1[0], p2[0], num_inpl)
            y_fit = np.array([np.polyval(coefficients, x) for x in nx]).T
            if len(point) == 3: 
                x, y, z = np.vstack((nx, y_fit))
            else:
                x, y = nx, y_fit
            
            if yMin is not None and yMax is not None:
                y = np.clip(y, yMin, yMax)
            if len(point) ==3 and zMin is not None and zMax is not None:
                z = np.clip(z, zMin, zMax)
        elif frame > 0:
            prev_point = trajectory[frame-1]
            current_point = point
            # 生成中间的点来模拟抛物线
            x = np.linspace(prev_point[0], current_point[0], num_inpl)
            y = np.linspace(prev_point[1], current_point[1], num_inpl)
            if len(point) ==3:
                z = np.linspace(prev_point[2], current_point[2], num_inpl)
    
    return x, y, z
    
def interpolate_trajectory(trajectory, intp_space=None):
    
    xMax, xMin = trajectory[:, 0].max(), trajectory[:, 0].min()
    yMax, yMin = trajectory[:, 1].max(), trajectory[:, 1].min()
    if len(trajectory[0, :]) == 3:
        zMax, zMin = trajectory[:, 2].max(), trajectory[:, 2].min()
    else:
        zMax, zMin = None, None
    
    interpolated_trajectory_3d = []
    for frame, point in enumerate(trajectory):
        
        x, y, z = cell_interpolate(frame, trajectory, inpl_space=intp_space, 
                    yMax=yMax, yMin=yMin, zMax=zMax, zMin=zMin)

        if x is not None and y is not None:
            if z is not None:
                interpolated_trajectory_3d.extend((np.vstack((x, y, z)).T)[1:-1])
            else:
                interpolated_trajectory_3d.extend((np.vstack((x, y)).T)[1:-1])
    
        interpolated_trajectory_3d.append(point)
        
    interpolated_trajectory_3d = np.array(interpolated_trajectory_3d)
    
    return interpolated_trajectory_3d

def draw_trajectory_3d_animation(point_trajectory_3d, inpl_space=None, save_gif=False, filename=None):
    
    assert point_trajectory_3d.shape[1] == 3, "point_trajectory_3d should be N x 3"
    
    fig = plt.figure(figsize=(12, 8))
    # ax = p3.Axes3D(fig)
    
    ax: Axes3D = plt.subplot(111, projection='3d')
    
    # 计算每个点到下一个点的方向向量
    directions = np.diff(point_trajectory_3d, axis=0)

    # 为了在每个点上绘制箭头，我们需要在每个点上加上一个起点
    start_points = point_trajectory_3d[:-1]

    # 初始化散点和箭头
    ax.scatter(point_trajectory_3d[0, 0], point_trajectory_3d[0, 1], point_trajectory_3d[0, 2], 
               color='green', marker='o', s=100, label='Start')
    scatter = ax.scatter(start_points[0, 0], start_points[0, 1], start_points[0, 2], c='r') # , label='trajectory'
    # quiver = ax.quiver(start_points[0, 0], start_points[0, 1], start_points[0, 2], 
    #                    directions[0, 0], directions[0, 1], directions[0, 2], 
    #                    color='y', length=0.5, label='direction')
    ax.scatter(point_trajectory_3d[-1, 0], point_trajectory_3d[-1, 1], point_trajectory_3d[-1, 2], 
               color='k', marker='*', s=100, label='End')
    # 设置坐标轴范围
    xMax, xMin = point_trajectory_3d[:, 0].max(), point_trajectory_3d[:, 0].min()
    yMax, yMin = point_trajectory_3d[:, 1].max(), point_trajectory_3d[:, 1].min()
    zMax, zMin = point_trajectory_3d[:, 2].max(), point_trajectory_3d[:, 2].min()
    xPadding = 0.2 * (xMax - xMin)    
    yPadding = 0.2 * (yMax - yMin)
    zPadding = 0.2 * (zMax - zMin)
    ax.set_xlim(xMin - xPadding, xMax + xPadding)
    ax.set_ylim(yMin - yPadding, yMax + yPadding)
    ax.set_zlim(zMin - zPadding, zMax + zPadding)

    ax.legend()
    ax.set_title('transform_trajectory')
    ax.set_xlabel('X (W)')
    ax.set_ylabel('Y (D)')
    ax.set_zlabel('Z (H)')
    ax.grid(True)

    # 更新箭头位置的函数
    def update(frame):
        # ax.plot(point_trajectory_3d[frame-1:frame+1, 0], 
        #         point_trajectory_3d[frame-1:frame+1, 1], 
        #         point_trajectory_3d[frame-1:frame+1, 2], 
        #         c='r', marker='o', ls='dotted', lw=1)
        # 更新点位置
        ax.scatter(point_trajectory_3d[frame, 0], 
                   point_trajectory_3d[frame, 1], 
                   point_trajectory_3d[frame, 2], 
                   c='r')
        
        # x, y, z = start_points[frame, :]
        # u, v, w = directions[frame, :]
        # ax.quiver(x, y, z, u, v, w, color='y', length=0.1, arrow_length_ratio=0.3, label='direction')
        
        # x, y, z = None, None, None
        # num_interpolation = 10
        # if frame > 1:
        #     # 获取前一个点和当前点
        #     p0 = point_trajectory_3d[frame-2]
        #     p1 = point_trajectory_3d[frame-1]
        #     p2 = point_trajectory_3d[frame]
            
        #     pp = np.vstack((p0, p1, p2))
        #     x, y, z = pp.T
            
        #     # 使用 np.polyfit 拟合一个二次函数
        #     yz = np.vstack((y, z)).T
        #     coefficients = np.polyfit(x, yz, 2)
            
        #     nx = np.linspace(p1[0], p2[0], num_interpolation)
        #     yz_fit = np.array([np.polyval(coefficients, x) for x in nx]).T
        #     x, y, z = np.vstack((nx, yz_fit))
        #     y = np.clip(y, yMin, yMax)
        #     z = np.clip(z, zMin, zMax)
        # elif frame > 0:
        #     prev_point = point_trajectory_3d[frame-1]
        #     current_point = point_trajectory_3d[frame]
        #     # 生成中间的点来模拟抛物线
        #     x = np.linspace(prev_point[0], current_point[0], num_interpolation)
        #     y = np.linspace(prev_point[1], current_point[1], num_interpolation)
        #     z = np.linspace(prev_point[2], current_point[2], num_interpolation)

        x, y, z = cell_interpolate(frame, point_trajectory_3d, inpl_space=inpl_space, 
                                   yMax=yMax, yMin=yMin, zMax=zMax, zMin=zMin)
        
        if x is not None and y is not None and z is not None:
            # 绘制抛物线连接的点
            ax.plot(x, y, z, c='y', marker='.', ls='dotted', lw=1)
        
    # 创建动画
    ani = FuncAnimation(fig, update, frames=len(directions), interval=100, repeat=False)
    
    if save_gif:
        if filename is None:
            # 保存动画为GIF, 标题增加时间戳
            import time
            timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
            filename = f'_3D_{timestamp}.gif'
        ani.save(filename, writer='pillow', fps=10)
        # ani.save('trajectory.gif', writer='pillow', fps=10)

    # 显示动画
    plt.show()
    

def test_3d_draw():
    fig = plt.figure()
    # ax = p3.Axes3D(fig)
    ax: Axes3D = plt.subplot(111, projection='3d')
    
    # 通过子图创建 Axes3D 对象
    # ax = plt.subplot(111, projection='3d')

    x = np.random.standard_normal(100)
    y = np.random.standard_normal(100)
    z = np.random.standard_normal(100)

    
    # x=np.arange(-8,8,0.25)     # 生成数据
    # y=np.arange(-8,8,0.25)
    # x,y=np.meshgrid(x,y)   #生成x,y轴数据
    # r=np.sqrt(x**2+y**2)
    # z=np.sin(r)/r                 #生成z值

    # 数据坐标
    z = np.linspace(0, 15, 1000)
    x = np.sin(z)
    y = np.cos(z)

    # 绘制线性图
    ax.plot(x, y, z)
    plt.show()

    
def test_draw_2d(traj_file_name = 'results_add_keypoints_track_724_trajectory.txt', anim=0, save_gif=1):
    # traj_file_name = 'results_add_keypoints_track_724_trajectory.txt'
    # traj_file_name = 'elbow_center_mask.txt'
    traj_file_name = os.path.join(WORK_DIR, traj_file_name)
    point_trajectory = np.loadtxt(traj_file_name)
    
    # # change the order of x and y
    # point_trajectory = point_trajectory[:, [1, 0]]
    
    # result = check_traj_continuity(point_trajectory)
    # point_jumps = result['point_jumps']
    point_jumps = None
    
    # point_trajectory = interpolate_trajectory(point_trajectory, intp_space=result.get("diff_mean"))
    
    if not anim:
        draw_trajectory_2d(point_trajectory, jumps=point_jumps, 
                       save_gif=save_gif, filename=traj_file_name.split('.')[0]+'_2D.jpg')
    else:
        draw_trajectory_2d_animation(point_trajectory, jumps=point_jumps, 
                                 save_gif=save_gif, filename=traj_file_name.split('.')[0]+'_2D.gif')
    

def test_3d_traj_draw(traj_file_name = 'results_add_keypoints_track_724_trajectory.txt', anim=0, save_fig=1):
    # x = np.random.standard_normal(100)
    # y = np.random.standard_normal(100)
    # z = np.random.standard_normal(100)
    
    # traj_file_name = 'results_add_keypoints_track_724_trajectory.txt'
    traj_file_name = os.path.join(WORK_DIR, traj_file_name)
    point_trajectory = np.loadtxt(traj_file_name)
    
    z = point_trajectory[:, 0] # H
    x = point_trajectory[:, 1] # W
    
    # x = point_trajectory[:, 0] # H
    # z = point_trajectory[:, 1] # W
    
     # y = np.random.standard_normal(len(point_trajectory)) * 255
    y = np.loadtxt('elbow/trj_depth.txt')
   
    traj_3d = np.vstack((x, y, z)).T
    
    result = check_traj_continuity(traj_3d)
    
    # traj_3d = interpolate_trajectory(traj_3d, intp_space=result.get("diff_mean"))
    if not anim:
        draw_trajectory_3d(traj_3d, save_fig=save_fig, filename=traj_file_name.split('.')[0]+'_3D.jpg')
    else:
        draw_trajectory_3d_animation(traj_3d, inpl_space=result.get("diff_mean")/2, 
                                 save_gif=save_fig, filename=traj_file_name.split('.')[0]+'_3D.gif')
    
    # fig = plt.figure()
    # ax = fig.add_subplot(111, projection='3d')

    # # x, y, z是我们采集到的三维数据
    # ax.scatter(x, y, z)
    # # 获取散点对象
    # scatter = ax.collections[0]

    # # 更新散点数据
    # scatter.set_offsets((x, y, z))

    # # 更新绘制
    # fig.canvas.draw()
    
    # # plt.show()


def test_3d_polyfit():

    # 随机生成一些三维数据
    x = np.random.normal(size=3)
    y = np.random.normal(size=3)
    z = x**2 + y**2 + np.random.normal(size=3)
    print(np.vstack((x, y, z)))
    
    # 使用 np.polyfit 拟合一个二次函数
    yz = np.vstack((y, z)).T
    coefficients = np.polyfit(x, yz, 2)
    # print(coefficients)

    nx = 2.5
    yz_fit = np.polyval(coefficients, nx)
    print(f"xyz: {np.hstack((nx, yz_fit))}")
    nx = [1.5, 2.5, 3.5]
    yz_fit = np.array([np.polyval(coefficients, x) for x in nx]).T
    new_xyz = np.vstack((nx, yz_fit)).T
    print(new_xyz)
    # print(np.vstack((nx, z_fit)))

    # # 绘制原始散点图和拟合曲面
    # fig = plt.figure()
    # ax.plot_surface(x, y, z_fit, alpha=0.5)
    # ax.scatter(x, y, z)
    # ax.plot_surface(x_grid, y_grid, z_fit, alpha=0.5)
    # plt.show()


WORK_DIR = os.path.dirname(os.path.abspath(__file__))
WORK_DIR = os.path.join(WORK_DIR, "elbow")

if __name__ == '__main__':
    traj_file_name = 'results_add_goe_feats_track_724_trajectory.txt'
    # traj_file_name = 'results_add_keypoints_track_724_trajectory.txt'
    # traj_file_name = 'results_only_optical_flow_track_724_trajectory_compensate.txt'
    # traj_file_name = 'results_add_goe_feats_track_724_trajectory_camara_compenated.txt'
    # traj_file_name = 'results_add_goe_feats_track_724_trajectory_camara_compenated_right_hand_compenated.txt'
    # test_3d_polyfit()
    test_draw_2d(traj_file_name, anim=0, save_gif=1)
    # test_3d_traj_draw(traj_file_name, anim=1, save_fig=1)
    # test_3d_draw()